"""
@author: chenzhenhua
@project: jf_fashion
@file: pruning.py
@time: 2021/8/3 0003 15:33
@desc:
"""

import numpy as np
import pandas as pd
from tensorflow.python.keras import backend as K
from tensorflow.python.keras.models import Model


# 获取权重和每层输出
def get_wb(model, x_train):
    layer_data = []
    weights = []
    biases = []
    for layer in model.layers:
        value_layer_model = Model(inputs=model.input, outputs=model.get_layer(layer.name).output)
        weight, bias = model.get_layer(layer.name).get_weights()
        weights.append(weight)
        biases.append(bias)
        layer_data.append(value_layer_model.predict(x_train))
    return layer_data, weights, biases


# 相关度计算
def calculate_corr(hidden_layer_data, corr_rate):
    dd = pd.DataFrame(hidden_layer_data)
    corr = dd.corr(method='pearson')

    data_np = np.array(corr)
    data_np[data_np == 1] = 0

    corr_index = []
    corri_index = []
    corr_sum = []
    for i in range(len(data_np)):
        corri_sum = 1
        for j in range(i + 1, len(data_np)):
            if data_np[i, j] > corr_rate:
                corr_index.append(j)
                corri_sum = corri_sum + data_np[i, j]
        if corri_sum != 1:
            corri_index.append(i)
            corr_sum.append(corri_sum)

    return corr_index, corri_index, corr_sum


def forward(x0, weights, biases):
    layer1_in = x0 @ weights[0] + biases[0]
    layer1_out = K.relu(layer1_in)
    layer2_in = layer1_out @ weights[1] + biases[1]
    layer2_out = K.relu(layer2_in)
    layer3_in = layer2_out @ weights[2] + biases[2]
    layer3_out = K.softmax(layer3_in)
    return K.eval(layer3_out)


def new_weight(weights, biases, layer_data,add_corr_value=False):
    corr_index, corri_index, corr_sum = calculate_corr(layer_data[0], 0.7)

    if add_corr_value:
        for i in range(len(corri_index)):
            weights[0][:, corri_index[i]] = weights[0][:, corri_index[i]] * corr_sum[i]
            weights[1][corri_index[i], :] = weights[1][corri_index[i]] * corr_sum[i]

    corr_index_use = list(set(range(0, 256)) - set(corr_index))

    weights[0] = np.delete(weights[0], corr_index, axis=1)
    biases[0] = biases[0][corr_index_use]

    weights[1] = np.delete(weights[1], corr_index, axis=0)

    return weights, biases
